# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/12 19:08 
@Author  : ZhangShenao 
@File    : 6.tokenizer.py 
@Desc    : 使用OpenAI的tiktoken库,计算token长度

不同的大模型,采用了不同的编码器,计算出的token长度也不同
"""
import tiktoken


def count_token(text: str, model_name: str = "gpt-4o-mini") -> int:
    """
    计算文本的token数量
    :param text: 文本
    :param model_name: 模型名称,默认为gpt-4o-mini
    :return: token数量
    """

    # 获取模型对应的token编码器
    encoder = tiktoken.encoding_for_model(model_name)

    # 对文本进行编码
    tokens = encoder.encode(text)

    # 返回token长度
    return len(tokens)


if __name__ == '__main__':
    # 计算token长度
    token_count = count_token(text="中国的首都是哪里？", model_name="gpt-4o-mini")
    print(token_count)  # 6
    token_count = count_token(text="中国的首都是哪里？", model_name="gpt-3.5-turbo")
    print(token_count)  # 9
